import pandas as pd
import sqlalchemy

from agent.utils.get_config import config_data
from agent.utils.llm_access.LLM import get_llm

DATABASE_URL = config_data['mysql']
engine = sqlalchemy.create_engine(DATABASE_URL)

STATIC_URL = config_data['static_path']

llm = get_llm()

from .prediction.grade_prediction import predict_for_stu_func, predict_for_class_func
from .copilot.sql_code import query_database_func
from .copilot.python_code import draw_graph_func
from .rag.rag import rag_from_policy_func
from .db.query_db import from_lesson_name_to_lesson_num_func, from_username_to_uid_func


def predict_grade_for_stu(user_id: str, lesson_num: str) -> (float, str):
    """
    predict_grade_for_stu(user_id: str, lesson_num: str) -> (float, str):
    Predicts the grades for a student with uid and lesson number.
    Returns a float of grade predicted and a string or image path.

    Args:
    - user_id (str): The uid of the student for which to predict grades.
    - lesson_num (str): The number or identifier of the lesson for which grades are to be predicted.

    Returns:
    - float: A float number of the predicted grades.
    - str: A string of the path of the prediction image
    return None, None if prediction failed.

    Example:
    ```python
        predicted_grades, img_path= predict_grade_for_stu('1214234', 'Lesson42343')

        # Output(float):
        # 96.50
    ```
    """
    result, img_path = predict_for_stu_func(user_id, lesson_num, engine)
    return result, STATIC_URL + img_path[2:]


def predict_grade_for_class(class_name: str, lesson_num: str) -> (pd.DataFrame, str):
    """
    predict_grade_for_class(class_name: str, lesson_num: str) -> (pd.DataFrame, str):
    Predicts the grades for students in a class with class name and lesson number.
    Returns a pandas DataFrame with detailed information and a string or image path..

    Args:
    - class_name (str): The name of the class for which to predict grades.
    - lesson_num (str): The number or identifier of the lesson for which grades are to be predicted.

    Returns:
    - pd.DataFrame: A DataFrame containing the predicted grades along with student information.
        The DataFrame includes the following columns:
        - 'uid' (str): Unique identifier of the student.
        - 'predicted_grade' (float): The predicted grade for the student for the specified lesson.
        - 'stu_num' (str): The student number.
        - 'username' (str): The username of the student.
    - str: A string of the path of the prediction image
    return None, None if prediction failed.

    If no grades are found for any student in the class for the pre-lessons of the specified lesson number,
    the function will return None.

    Example:
    ```python
        predicted_grades_df, img_path = predict_grade_for_class('ClassA', 'Lesson3')

        # Output1(pd.DataFrame):
        #    uid  predicted_grade stu_num   username
        # 0  uid1           85.0   S001  John Doe
        # 1  uid2           92.5   S002  Jane Smith
        # 2  uid3           78.0   S003  Alice Johnson
        # ... and so on for each student in the class

        # Output2(str):
        # http://127.0.0.1:8001/tmp_imgs/cqwzdbko.png
    ```
    """
    result, img_path = predict_for_class_func(class_name, lesson_num, engine)
    result = pd.DataFrame(result)
    return result, STATIC_URL + img_path[2:]


def query_database(question: str, df_cols: str | list = None) -> pd.DataFrame:
    """
    query_database(question: str, df_cols: str | list = None) -> pd.DataFrame:
    Query the database using natural language question. Can not query anything not included in the database content!!!
    Returns the query results in pandas DataFrame.

    Args:
    - question (str): Natural language question.
    - df_cols (str|list): The columns' names of the DataFrame(e.g. "uid, username, stu_num").

    Returns:
    - pd.DataFrame: A DataFrame containing the results of the database query.
        The DataFrame includes the columns provided in df_cols(the second args)

    Example:
    ```python
        ans_df = query_database('Select the grades of Jane Smith', df_cols='lesson_id, lesson_name, grade')

        # Output1(pd.DataFrame):
        #        lesson_id lesson_name grade
        # 0        001  Mathematics     99.00
        # 1        002      English     88.50
        # 2        003     Physics    65.00
        # ... and so on(the structure of the output DataFrame id based on df_cols(the second input args))

        # Output2(str):
        # http://127.0.0.1:8001/tmp_imgs/cqwzdbko.png
    ```
    """

    result = query_database_func(question, df_cols, llm, engine)
    return result


def draw_graph(question: str, data: pd.DataFrame) -> str:
    """
    draw_graph(question: str, data: pd.DataFrame) -> str:
    Draw graph based on natural language graph type and data provided in a pandas DataFrame.
    Returns an url path string of the graph.

    Args:
    - question (str): Natural language graph type.
    - data (pd.DataFrame): A pandas DataFrame for providing drawing data.

    Returns:
    - str: url path string of the output graph.(e.g. "http://127.0.0.1:8003/tmp_imgs/mlkjcvep.png").

    Example:
    ```python
        data = pd.DataFrame({
            '月份': ['1月', '2月', '3月', '4月', '5月'],
            '销售额': [200, 220, 250, 210, 230]
        })
        graph_url = draw_graph("画折线图", data)
        # Output(str):
        # "http://127.0.0.1:8003/tmp_imgs/ekhidpcl.png"
    ```
    """
    result = draw_graph_func(question, data, llm)
    result = STATIC_URL + result[2:]
    return result


def rag_from_policy(question: str) -> str:
    """
    rag_from_policy(question: str) -> str:
    Retrieve and Augment Generation (RAG) from school policy knowledge base using natural language question.
    Returns a string of the answer.

    Args:
    - question (str): Natural language question related to school policy.

    Returns:
    - str: A coherent and informative response based on the question and the school policy data.

    Example:
    ```python
        response = rag_from_policy("What is the school's policy on graduation?")
        # Output(str):
        # "According to the school's policy, ....."
    ```
    """
    result = rag_from_policy_func(question, llm, engine)
    return result


def from_lesson_name_to_lesson_num(lesson_name: str) -> str:
    """
    from_lesson_name_to_lesson_num(lesson_name: str) -> str:
    Convert lesson name(lesson_name) to lesson number(lesson_num)
    Returns a string of lesson number.

    Args:
    - lesson_name (str): The Name of a Lesson.

    Returns:
    - str: The lesson number(lesson_num) of the Lesson.
    return None if no lesson number is found with the lesson name
    """
    result = from_lesson_name_to_lesson_num_func(lesson_name, engine)
    return result


def from_username_to_uid(username: str) -> str:
    """
    from_username_to_uid(username: str) -> str:
    Convert username to uid
    Returns a string of uid.

    Args:
    - username (str): The Name of a user(username).

    Returns:
    - str: The uid of the user.
    return None if no user is found with the username
    """
    result = from_username_to_uid_func(username, engine)
    return result
